import datetime
import logging
import os
import time
import ujson
from absl import app, flags
from collections import namedtuple
import pandas as pd
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from scipy.stats import norm

import coin2.service.strategy.trading_summary_service_pb2 as tss_pb2
from xunkemgmt_client.base.datetime_util import to_timestamp_int, to_datetime
from xunkemgmt_client.client.api_client import XunkemgmtClient
from python.experimental.lucasyoo.accounting.offset_accounting2 import (Data)
import python.experimental.lucasyoo.accounting.accounting as lac

matplotlib.use("Agg")
os.environ["TZ"] = "GMT"
time.tzset()

FLAGS = flags.FLAGS

CORP_OPTION_PNL = namedtuple('CORP_OPTION_PNL', ['base', 'trading_date', 'otc_pnl_in_usd'])


def plot_and_save(df: pd.DataFrame, base: str, output_dir: str):
  start_dt = df.index.min().strftime("%Y%m%d")
  end_dt = df.index.max().strftime("%Y%m%d")
  df.to_csv(f'{output_dir}/{base}_{start_dt}_{end_dt}.csv')
  plt.rcParams['figure.figsize'] = (30, 20)
  plt.plot(df.index,
           df['otc_pnl_in_usd'],
           'b-',
           lw=0.5,
           drawstyle='steps-post',
           label='otc_pnl_in_usd')
  plt.legend()
  plt.title(f"{base} ({start_dt}-{end_dt})"
            f"(option: {df['otc_pnl_in_usd'].iloc[-1]:,.0f})")
  plt.savefig(f'{output_dir}/{base}_{start_dt}_{end_dt}.png')
  plt.close()


def plot_total(start_time: datetime.datetime, end_time: datetime.datetime, output_dir: str):
  strat_base = []
  tmp_dfs = []
  with XunkemgmtClient() as client:
    rsp = client.query_trading_summary_interval_histories(
        tss_pb2.QueryTradingSummaryIntervalHistoriesRequestProto(
            business_units=['Day1mm'],
            strategy_groups=['vmm'],
            start_timestamp=to_timestamp_int(start_time),
            end_timestamp=to_timestamp_int(end_time),
            agg_types=['STRATEGY_NAME', 'TRADING_DATE']))
    for his in rsp.histories:
      strategy_name = his.summary_info.strategy_name
      if "delta_hedge" not in strategy_name:
        continue
      base = strategy_name.replace("vmm_", "").replace("_delta_hedge", "")
      if "_" in base:
        base = base.split("_")[0]
      base = base.upper()
      if FLAGS.symbol != "ALL" and FLAGS.symbol != base:
        continue
      strat_base.append((strategy_name, base, str(his.summary_info.trading_date)))
  strat_info = pd.DataFrame(strat_base, columns=["strategy_name", "base", "trading_date"])
  strat_info['trading_date'] = pd.to_datetime(strat_info['trading_date'])
  unique_strat_info = strat_info[["strategy_name", "base"]].drop_duplicates()
  for _, (strategy_name, base) in unique_strat_info.iterrows():
    tmp_start_time = strat_info.loc[(strat_info['strategy_name'] == strategy_name) &
                                    (strat_info['base'] == base), 'trading_date'].min()
    # tmp_end_time = strat_info.loc[(strat_info['strategy_name'] == strategy_name) &
    #                               (strat_info['base'] == base), 'trading_date'].max()
    tmp_df, _, _ = lac.get_option_pnl(tmp_start_time - datetime.timedelta(days=1), end_time,
                                      strategy_name, base)
    tmp_dfs.append(tmp_df)
  df = pd.concat(tmp_dfs)
  if df.empty:
    return
  sliced_df = df[['ts', 'base', 'call_price_diff']].copy()
  config = lac.query_option_num_info()
  sliced_df = sliced_df.merge(pd.DataFrame(config).T.reset_index(names="base"),
                              on="base",
                              how='left')
  sliced_df["otc_pnl_in_usd"] = sliced_df["call_price_diff"] * (sliced_df["total"] -
                                                                sliced_df['team_flow'])
  sliced_df.rename(columns={"ts": "trading_date"}, inplace=True)
  total_pnl = sliced_df.groupby(["trading_date"])[["otc_pnl_in_usd"]].sum().cumsum()
  plot_and_save(total_pnl, 'TOTAL', output_dir)
  for base in sliced_df["base"].unique():
    tmp = sliced_df[sliced_df['base'] == base]
    tmp = tmp.groupby(["trading_date"])[["otc_pnl_in_usd"]].sum().cumsum()
    plot_and_save(tmp, base, output_dir)


def main(_):
  start_time = datetime.datetime.strptime(FLAGS.start_time, '%Y%m%d')
  end_time = datetime.datetime.strptime(FLAGS.end_time, '%Y%m%d')
  output_dir = os.path.join(os.path.expandvars('$COIN_REPO'), FLAGS.output_dir)
  os.makedirs(output_dir, exist_ok=True)
  plot_total(start_time, end_time, output_dir)


if __name__ == "__main__":
  flags.DEFINE_string('start_time', '2021-01-01', 'start time')
  flags.DEFINE_string('end_time', '2021-01-02', 'end time')
  flags.DEFINE_string('output_dir', 'sim_result/accounting', 'output directory')
  flags.DEFINE_string('symbol', 'ALL', 'symbol')
  app.run(main)
